#pragma once

#include "concepts.h"
#include "task.h"
#include <cerrno>
#include <chrono>
#include <coroutine>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <functional>
#include <liburing.h>
#include <new>
#include <span>
#include <utility>

namespace co_async {
using UringTask = Task<int>;

template <class Rep, class Period>
inline __kernel_timespec
duration_to_kernel_timespec(const std::chrono::duration<Rep, Period> &dur) {
  auto sec = std::chrono::duration_cast<std::chrono::seconds>(dur);
  return (struct __kernel_timespec){
      .tv_sec = sec.count(),
      .tv_nsec = std::chrono::duration_cast<std::chrono::nanoseconds>(dur - sec)
                     .count(),
  };
}

template <class Clk, class Dur>
inline __kernel_timespec
time_point_to_kernel_timespec(std::chrono::time_point<Clk, Dur> tp) {
  return duration_to_kernel_timespec(tp.time_since_epoch());
}

struct UringLoop {
  io_uring m_ring;
  std::deque<std::coroutine_handle<>> m_handles;
  io_uring_sqe *get_sqe() {
    io_uring_sqe *res = io_uring_get_sqe(&m_ring);
    if (!res) [[unlikely]] {
      throw std::bad_alloc();
    }
    return res;
  }
  explicit UringLoop(std::size_t entries = 512) {
    check_uring_error(io_uring_queue_init(entries, &m_ring, 0));
  }

  inline void run();
  UringLoop &operator=(UringLoop &&) = delete;
  ~UringLoop() { io_uring_queue_exit(&m_ring); }
};
struct UringAwaiter {
  explicit UringAwaiter(UringLoop &loop, const auto &func) : m_loop(loop) {
    io_uring_sqe *sqe = m_loop.get_sqe();
    func(sqe); // 设置sqe任务
    io_uring_sqe_set_data(sqe, this);
  }
  bool await_ready() const noexcept { return false; }
  void await_suspend(std::coroutine_handle<> handle) {
    m_pre = handle;
    check_uring_error(io_uring_submit(&m_loop.m_ring));
  }
  int await_resume() const noexcept { return m_res; }
  UringLoop &m_loop;
  std::coroutine_handle<> m_pre; // 回调任务
  int m_res = -ENOSYS;
};

void UringLoop::run() {
  while (!m_handles.empty()) {
    auto handle = std::move(m_handles.front());
    m_handles.pop_front();
    handle.resume();
  }
  io_uring_cqe *cqe;
  check_uring_error(io_uring_wait_cqe(&m_ring, &cqe));
  auto awaiter = reinterpret_cast<UringAwaiter *>(cqe->user_data);
  awaiter->m_res = cqe->res;
  io_uring_cqe_seen(&m_ring, cqe);
  awaiter->m_pre.resume();
}

// 使用指定的文件描述符和路径打开一个文件
inline UringTask uring_openat(UringLoop &loop, int dfd, const char *path,
                              int flags, mode_t mode = 0644) {
  co_return check_uring_error(
      co_await UringAwaiter{loop, [&](auto sqe) {
                              io_uring_prep_openat(sqe, dfd, path, flags, mode);
                            }});
}

inline UringTask uring_socket(UringLoop &loop, int domain, int type,
                              int protocol, unsigned int flags) {
  co_return check_uring_error(co_await UringAwaiter{
      loop, [&](auto sqe) {
        io_uring_prep_socket(sqe, domain, type, protocol, flags);
      }});
}

// 接受一个连接
inline UringTask uring_accept(UringLoop &loop, int fd, struct sockaddr *addr,
                              socklen_t *addrlen, int flags) {
  co_return check_uring_error(co_await UringAwaiter{
      loop,
      [&](auto sqe) { io_uring_prep_accept(sqe, fd, addr, addrlen, flags); }});
}

// 连接一个套接字
inline Task<int> uring_connect(UringLoop &loop, int fd,
                               const struct sockaddr *addr, socklen_t addrlen) {
  co_return check_uring_error(co_await UringAwaiter{
      loop, [&](auto sqe) { io_uring_prep_connect(sqe, fd, addr, addrlen); }});
}

// 创建一个目录
inline Task<int> uring_mkdirat(UringLoop &loop, int dfd, char const *path,
                               mode_t mode = 0644) {
  co_return check_uring_error(co_await UringAwaiter{
      loop, [&](auto sqe) { io_uring_prep_mkdirat(sqe, dfd, path, mode); }});
}

// 创建一个硬链接
inline Task<int> uring_linkat(UringLoop &loop, int olddfd, char const *oldpath,
                              int newdfd, char const *newpath, int flags) {

  co_return check_uring_error(co_await UringAwaiter(loop, [&](auto sqe) {
    io_uring_prep_linkat(sqe, olddfd, oldpath, newdfd, newpath, flags);
  }));
}

// 删除一个链接
inline Task<int> uring_unlinkat(UringLoop &loop, int dirfd, char const *path,
                                int flags = 0) {

  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_unlinkat(sqe, dirfd, path, flags);
      }));
}

// 创建一个符号链接
inline Task<int> uring_symlinkat(UringLoop &loop, char const *target,
                                 int newdirfd, char const *linkpath) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_symlinkat(sqe, target, newdirfd, linkpath);
      }));
}

// 获取文件状态
inline Task<int> uring_statx(UringLoop &loop, int dirfd, char const *path,
                             int flags, unsigned int mask,
                             struct statx *statxbuf) {
  co_return co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
    io_uring_prep_statx(sqe, dirfd, path, flags, mask, statxbuf);
  });
}

inline Task<std::size_t> uring_read(UringLoop &loop, int fd,
                                    std::span<char> buf,
                                    std::uint64_t offset = 0) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_read(sqe, fd, buf.data(), buf.size(), offset);
      }));
}

inline Task<std::size_t> uring_write(UringLoop &loop, int fd,
                                     std::span<char const> buf,
                                     std::uint64_t offset = 0) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_write(sqe, fd, buf.data(), buf.size(), offset);
      }));
}

inline Task<std::size_t> uring_readv(UringLoop &loop, int fd,
                                     std::span<iovec const> buf,
                                     std::uint64_t offset, int flags) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_readv2(sqe, fd, buf.data(), buf.size(), offset, flags);
      }));
}

inline Task<std::size_t> uring_writev(UringLoop &loop, int fd,
                                      std::span<iovec const> buf,
                                      std::uint64_t offset, int flags) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_writev2(sqe, fd, buf.data(), buf.size(), offset, flags);
      }));
}

inline Task<std::size_t> uring_recv(UringLoop &loop, int fd,
                                    std::span<char> buf, int flags) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_recv(sqe, fd, buf.data(), buf.size(), flags);
      }));
}

inline Task<std::size_t> uring_send(UringLoop &loop, int fd,
                                    std::span<char const> buf, int flags) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_send(sqe, fd, buf.data(), buf.size(), flags);
      }));
}

inline Task<std::size_t> uring_recvmsg(UringLoop &loop, int fd,
                                       struct msghdr *msg, unsigned int flags) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_recvmsg(sqe, fd, msg, flags);
      }));
}

inline Task<std::size_t> uring_sendmsg(UringLoop &loop, int fd,
                                       struct msghdr *msg, unsigned int flags) {
  co_return check_uring_error(
      co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
        io_uring_prep_sendmsg(sqe, fd, msg, flags);
      }));
}

inline Task<int> uring_close(UringLoop &loop, int fd) {
  co_return check_uring_error(co_await UringAwaiter(
      loop, [&](io_uring_sqe *sqe) { io_uring_prep_close(sqe, fd); }));
}

inline Task<int> uring_shutdown(UringLoop &loop, int fd, int how) {
  co_return check_uring_error(co_await UringAwaiter(
      loop, [&](io_uring_sqe *sqe) { io_uring_prep_shutdown(sqe, fd, how); }));
}

inline Task<int> uring_fsync(UringLoop &loop, int fd, unsigned int flags) {
  co_return check_uring_error(co_await UringAwaiter(
      loop, [&](io_uring_sqe *sqe) { io_uring_prep_fsync(sqe, fd, flags); }));
}

inline Task<int> uring_cancel_fd(UringLoop &loop, int fd, unsigned int flags) {
  co_return co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
    io_uring_prep_cancel(sqe, nullptr, flags);
  });
}

inline Task<int> uring_timeout(UringLoop &loop, struct __kernel_timespec ts,
                               unsigned int count, unsigned int flags) {
  int res = co_await UringAwaiter(loop, [&](io_uring_sqe *sqe) {
    io_uring_prep_timeout(sqe, &ts, count, flags);
  });
  if (res == -ETIME) [[likely]]
    res = 0;
  co_return check_uring_error(res);
}

template <class Rep, class Period>
inline Task<int> uring_timeout(UringLoop &loop,
                               std::chrono::duration<Rep, Period> dur,
                               std::size_t count) {
  return uring_timeout(loop, duration_to_kernel_timespec(dur), count,
                       IORING_TIMEOUT_ETIME_SUCCESS | IORING_TIMEOUT_REALTIME);
}

template <class Clk, class Dur>
inline Task<int> uring_timeout(UringLoop &loop,
                               std::chrono::time_point<Clk, Dur> tp,
                               std::size_t count) {
  return uring_timeout(loop, time_point_to_kernel_timespec(tp), count,
                       IORING_TIMEOUT_ETIME_SUCCESS | IORING_TIMEOUT_ABS |
                           IORING_TIMEOUT_REALTIME);
}
template <class A>
  requires(Awaitable<A>)
inline void uring_enqueue(UringLoop &loop, A &&task) {
  auto coroutine =
      task.operator co_await().await_suspend(std::noop_coroutine());
  loop.m_handles.push_back(coroutine);
  task.m_handle = nullptr;
}

template <class A>
  requires(Awaitable<A>)
inline void uring_enqueue(UringLoop &loop, std::function<A()> task) {
  uring_enqueue(loop, task());
}

} // namespace co_async